import asyncio
import functools
import logging
import socket
import warnings
from types import TracebackType
from typing import Any
from typing import Awaitable
from typing import Generator
from typing import List
from typing import Optional
from typing import Set
from typing import Type

from channel.protocol import Protocol

logger = logging.getLogger(__name__)


class ServerProtocol(Protocol):
    side = "server"

    def __init__(self, handler, server=None):
        self.protocol_handler = handler
        # self.protocol_server = server
        super().__init__()

    def connection_made(self, transport: asyncio.BaseTransport) -> None:
        """
        Register connection and initialize a task to handle it.

        """
        super().connection_made(transport)
        # Register the connection with the server before creating the handler
        # task. Registering at the beginning of the handler coroutine would
        # create a race condition between the creation of the task, which
        # schedules its execution, and the moment the handler starts running.
        # self.protocol_server.register(self)
        self.handler_task = self.loop.create_task(self.handler())

    async def handler(self) -> None:
        """
        Handle the lifecycle of a WebSocket connection.

        Since this method doesn't have a caller able to handle exceptions, it
        attemps to log relevant ones and guarantees that the TCP connection is
        closed before exiting.

        """
        try:
            try:
                await self.handshake()
            except ConnectionError:
                logger.debug("Connection error in opening handshake", exc_info=True)
                raise
            except Exception as exc:
                logger.warning(f"connection {exc}")
                self.fail_connection()
                await self.wait_closed()
                return

            try:
                await self.protocol_handler(self)
            except Exception:
                logger.error("Error in connection handler", exc_info=True)
                raise

            try:
                await self.close()
            except ConnectionError:
                logger.debug("Connection error in closing handshake", exc_info=True)
                raise
            except Exception:
                logger.warning("Error in closing handshake", exc_info=True)
                raise

        except Exception:
            # Last-ditch attempt to avoid leaking connections on errors.
            try:
                self.transport.close()
            except Exception:  # pragma: no cover
                pass

        finally:
            # Unregister the connection with the server when the handler task
            # terminates. Registration is tied to the lifecycle of the handler
            # task because the server waits for tasks attached to registered
            # connections before terminating.
            # self.protocol_server.unregister(self)
            pass


class ProtocolServer(object):
    def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
        # Store a reference to loop to avoid relying on self.server._loop.
        self.loop = loop

        # Keep track of active connections.
        self.websockets: Set[ServerProtocol] = set()

        # Task responsible for closing the server and terminating connections.
        self.close_task: Optional[asyncio.Task[None]] = None

        # Completed when the server is closed and connections are terminated.
        self.closed_waiter: asyncio.Future[None] = loop.create_future()

    def wrap(self, server: asyncio.AbstractServer) -> None:
        """
        Attach to a given :class:`~asyncio.Server`.

        Since :meth:`~asyncio.loop.create_server` doesn't support injecting a
        custom ``Server`` class, the easiest solution that doesn't rely on
        private :mod:`asyncio` APIs is to:

        - instantiate a :class:`WebSocketServer`
        - give the protocol factory a reference to that instance
        - call :meth:`~asyncio.loop.create_server` with the factory
        - attach the resulting :class:`~asyncio.Server` with this method

        """
        self.server = server

    def register(self, protocol: ServerProtocol) -> None:
        """
        Register a connection with this server.

        """
        self.websockets.add(protocol)

    def unregister(self, protocol: ServerProtocol) -> None:
        """
        Unregister a connection with this server.

        """
        self.websockets.remove(protocol)

    def is_serving(self) -> bool:
        """
        Tell whether the server is accepting new connections or shutting down.

        """
        try:
            # Python ≥ 3.7
            return self.server.is_serving()
        except AttributeError:  # pragma: no cover
            # Python < 3.7
            return self.server.sockets is not None

    def close(self) -> None:
        """
        Close the server.

        This method:

        * closes the underlying :class:`~asyncio.Server`;
        * rejects new WebSocket connections with an HTTP 503 (service
          unavailable) error; this happens when the server accepted the TCP
          connection but didn't complete the WebSocket opening handshake prior
          to closing;
        * closes open WebSocket connections with close code 1001 (going away).

        :meth:`close` is idempotent.

        """
        if self.close_task is None:
            self.close_task = self.loop.create_task(self._close())

    async def _close(self) -> None:
        """
        Implementation of :meth:`close`.

        This calls :meth:`~asyncio.Server.close` on the underlying
        :class:`~asyncio.Server` object to stop accepting new connections and
        then closes open connections with close code 1001.

        """
        # Stop accepting new connections.
        self.server.close()

        # Wait until self.server.close() completes.
        await self.server.wait_closed()

        # Wait until all accepted connections reach connection_made() and call
        # register(). See https://bugs.python.org/issue34852 for details.
        await asyncio.sleep(0)

        # Close OPEN connections with status code 1001. Since the server was
        # closed, handshake() closes OPENING conections with a HTTP 503 error.
        # Wait until all connections are closed.

        # asyncio.wait doesn't accept an empty first argument
        if self.websockets:
            await asyncio.wait([websocket.close(1001) for websocket in self.websockets])

        # Wait until all connection handlers are complete.

        # asyncio.wait doesn't accept an empty first argument.
        if self.websockets:
            await asyncio.wait(
                [websocket.handler_task for websocket in self.websockets]
            )

        # Tell wait_closed() to return.
        self.closed_waiter.set_result(None)

    async def wait_closed(self) -> None:
        """
        Wait until the server is closed.

        When :meth:`wait_closed` returns, all TCP connections are closed and
        all connection handlers have returned.

        """
        await asyncio.shield(self.closed_waiter)

    @property
    def sockets(self) -> Optional[List[socket.socket]]:
        """
        List of :class:`~socket.socket` objects the server is listening to.

        ``None`` if the server is closed.

        """
        return self.server.sockets


class Serve:
    def __init__(self, protocol_handler, host, port):
        loop = asyncio.get_event_loop()
        protocol_server = ProtocolServer(loop)
        protocol_factory = ServerProtocol

        protocol_factory = functools.partial(
            protocol_factory, protocol_handler, protocol_server
        )

        self._create_server = functools.partial(
            loop.create_server, protocol_factory, host, port
        )
        self._protocol_server = protocol_server

    async def __aenter__(self) -> ProtocolServer:
        return await self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None:
        self._protocol_server.close()
        await self._protocol_server.wait_closed()

    # await serve(...)

    def __await__(self) -> Generator[Any, None, ProtocolServer]:
        # Create a suitable iterator by calling __await__ on a coroutine.
        return self.__await_impl().__await__()

    async def __await_impl(self) -> ProtocolServer:
        server = await self._create_server()
        self._protocol_server.wrap(server)
        return self._protocol_server


serve = Serve
